use alloc::{collections::btree_map::BTreeMap, vec::Vec};
use spin::{Lazy, Mutex};
use x86_64::{
    structures::paging::{Mapper, Page, PageTableFlags, PhysFrame, Size4KiB},
    PhysAddr, VirtAddr,
};

use crate::arch::{kmm::FRAME_ALLOCATOR, syscall::SysError};

use super::{SysResult, Syscall};

static C_ALLOCATION_MAP: Lazy<Mutex<BTreeMap<VirtAddr, (VirtAddr, usize, usize)>>> =
    Lazy::new(|| Mutex::new(BTreeMap::new()));

pub fn sbrk(size: usize) -> SysResult {
    let size = size as isize;
    if size > 0 {
        let size = size as usize;

        let space: Vec<u8> = alloc::vec![0u8; size];

        let (ptr, len, cap) = space.into_raw_parts();
        assert!(len == size);
        if !ptr.is_null() {
            let vaddr = VirtAddr::new(ptr as u64);
            let mut guard = C_ALLOCATION_MAP.lock();
            if guard.contains_key(&vaddr) {
                drop(guard);
                unsafe {
                    drop(Vec::from_raw_parts(vaddr.as_mut_ptr() as *mut u8, len, cap));
                }
                panic!(
                    "sbrk(): vaddr {:?} already exists in C Allocation Map, query size: {size}",
                    vaddr
                );
            }
            guard.insert(vaddr, (vaddr, len, cap));
            return Ok(vaddr.as_u64() as usize);
        } else {
            return Err(SysError::ENOMEM);
        }
    } else {
        return Err(SysError::EINVAL);
    }
}

impl Syscall<'_> {
    pub fn sys_mmap(&mut self, addr: usize, len: usize, _prot: usize, flags: usize) -> SysResult {
        if addr == 0 {
            return sbrk(len);
        }
        if flags == 0 {
            return Err(SysError::EINVAL);
        }

        let mut offset_page_table = self.thread().get_context().get_current_page_table();

        for i in 0..((len + 4095) / 4096) {
            let result = unsafe {
                offset_page_table.map_to(
                    Page::<Size4KiB>::containing_address(VirtAddr::new((addr + i * 4096) as u64)),
                    PhysFrame::containing_address(PhysAddr::new((addr + i * 4096) as u64)),
                    PageTableFlags::from_bits(flags as u64).unwrap(),
                    &mut *FRAME_ALLOCATOR.lock(),
                )
            };

            match result {
                Ok(f) => f.flush(),
                Err(e) => match e {
                    x86_64::structures::paging::mapper::MapToError::PageAlreadyMapped(_) => {
                        offset_page_table
                            .unmap(Page::<Size4KiB>::containing_address(VirtAddr::new(
                                (addr + i * 4096) as u64,
                            )))
                            .unwrap()
                            .1
                            .flush();
                        unsafe {
                            offset_page_table
                                .map_to(
                                    Page::<Size4KiB>::containing_address(VirtAddr::new(
                                        (addr + i * 4096) as u64,
                                    )),
                                    PhysFrame::containing_address(PhysAddr::new(
                                        (addr + i * 4096) as u64,
                                    )),
                                    PageTableFlags::from_bits(flags as u64).unwrap(),
                                    &mut *FRAME_ALLOCATOR.lock(),
                                )
                                .unwrap()
                                .flush();
                        }
                    }
                    _ => panic!("Cannot do mmap"),
                },
            }
        }

        Ok(addr)
    }

    pub fn sys_alloc_dma(&mut self, size: usize) -> SysResult {
        let addr = FRAME_ALLOCATOR
            .lock()
            .allocate_frames(size)
            .ok_or(SysError::ENOMEM)?;
        Ok(addr.start_address().as_u64() as usize)
    }

    pub fn sys_free_dma(&mut self, addr: usize, size: usize) -> SysResult {
        let addr = PhysAddr::new(addr as u64);
        FRAME_ALLOCATOR.lock().deallocate_frames(addr, size);
        Ok(0)
    }
}
